package com.ant.backstage.utils;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.*;

/**
 * 自定义子类，核心点在于自定义header的map，再重写获取header的方法，将自定义的map也一起返回
 * 重写HttpServletRequest的请求头
*/
public class HeaderMapRequestWrapper extends HttpServletRequestWrapper {

    private static Logger logger = LogManager.getLogger(HeaderMapRequestWrapper.class);


    // 用于存放自定义header的map
    private Map<String, String> headerMap = new HashMap<>();

    /**
     * 构造方法，可保持原有的request内容不变
     *
     * @param request
     */
    public HeaderMapRequestWrapper(HttpServletRequest request) {
        super(request);
        this.headerMap = new HashMap<>();
    }

    /**
     * 在header中添加key-value
     */
    public void addHeader(String key, String value) {
        headerMap.put(key, value);
    }

    @Override
    public String getHeader(String name) {
        logger.info("获取getHeader --->{}",name);
        String headerValue = super.getHeader(name);
        if (headerMap.containsKey(name)) {
            headerValue = headerMap.get(name);
        }
        return headerValue;
    }

    /**
     * 需要重写的方法，即获取请求头Key时，加上自定义请求头内容的key
     * get the Header names
     */
    @Override
    public Enumeration<String> getHeaderNames() {
        List<String> names = Collections.list(super.getHeaderNames());
        for (String name : headerMap.keySet()) {
            names.add(name);
        }
        return Collections.enumeration(names);
    }

    /**
     * 需要重写的方法，获取所有的header参数时，添加上自定义header的map
     * get the Header key-value
     */
    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> values = Collections.list(super.getHeaders(name));
        if (headerMap.containsKey(name)) {
            values = Arrays.asList(headerMap.get(name));
        }
        return Collections.enumeration(values);
    }

}
